"""张量的索引与数据筛选"""
import torch

a = torch.rand(4, 4)
print("a")
print(a)
print()

b = torch.rand(4, 4)
print("b")
print(b)
print()


def test_where():
    print("==== torch.where ====")
    a = torch.rand(4, 4)
    print("a")
    print(a)
    print()

    b = torch.rand(4, 4)
    print("b")
    print(b)
    print()

    print("torch.where(a > 0.5, a, b)")
    out = torch.where(a > 0.5, a, b)
    print(out)
    print()


def test_gather_and_index_select():
    print("==== torch.index_select ====")
    a = torch.rand(4, 4)
    print("a")
    print(a)
    print()

    """index_select输入的索引为行/列的索引（也就是向量的索引）"""
    print("torch.index_select(a, dim=0, index=torch.tensor([0, 3, 2]))")
    out = torch.index_select(a, dim=0, index=torch.tensor([0, 3, 2]))
    print(out)
    print(out.shape)
    print()

    print("==== torch.gather ====")
    a = torch.linspace(1, 16, 16).view(4, 4)
    print("a")
    print(a)
    print()

    """gather输入的索引为元素的索引"""
    print("""torch.gather(a, dim=0, index=torch.tensor([[0, 1, 1, 1],
                                                     [0, 1, 2, 2],
                                                     [0, 1, 3, 3]]))""")
    out = torch.gather(a, dim=0, index=torch.tensor([[0, 1, 1, 1],
                                                     [0, 1, 2, 2],
                                                     [0, 1, 3, 3]]))
    # dim=0, out[i, j, k] = input[index[i,j,k], j, k]
    # dim=1, out[i, j, k] = input[i, index[i,j,k], k]
    # dim=2, out[i, j, k] = input[i, j, index[i,j,k]]
    print(out)
    print(out.shape)
    print()


def test_masked_select():
    print("==== torch.masked_select ====")
    a = torch.linspace(1, 16, 16).view(4, 4)
    print("a")
    print(a)
    print()

    print("mask = torch.gt(a, 8)")
    mask = torch.gt(a, 8)
    print(mask)
    print()

    print("torch.masked_select(a, mask)")
    out = torch.masked_select(a, mask)
    print(out)
    print()


def test_take():
    """将高维的tensor当成向量，take输入的索引从0~元素的个数"""
    print("==== torch.take ====")
    a = torch.linspace(1, 16, 16).view(4, 4)
    print("a")
    print(a)
    print()

    print("torch.take(a, index=torch.tensor([0, 15, 13]))")
    out = torch.take(a, index=torch.tensor([0, 15, 13]))
    print(out)
    print()


def test_nonzero():
    print("==== torch.nonzero ====")
    a = torch.tensor([[0, 1, 2, 0], [2, 3, 0, 1]])
    print("a")
    print(a)
    print()

    # 输出的是非零元素的索引  -->   稀疏表示
    print("torch.nonzero(a)")
    out = torch.nonzero(a)
    print(out)
    print()


if __name__ == '__main__':
    test_where()
    test_gather_and_index_select()
    test_masked_select()
    test_take()
    test_nonzero()
